
#include "mat_mul_sub_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/tiling_api.h"
#include <cstdint>
using namespace matmul_tiling;

inline void _configC(uint32_t &nCore, int &baseM, int &baseN, int M, int N) {
  int coresN = 4;
  int coresM = 4;
  // 先确定N轴方向的cores数量
  while (baseN * coresN > N && baseN > 16)
    baseN /= 2;
  while (baseN * coresN > N && coresN > 2)
    coresN /= 2;
  while (baseM * coresM > M && baseM > 16)
    baseM /= 2;
  while (baseM * coresM > M && coresM > 2)
    coresM /= 2;
  nCore = coresM * coresN;
  if (nCore % 2)
    printf("[ERROR] Core Num invalid.\n");
}

inline void configCores(uint32_t &nCore, int &baseM, int &baseN, int M, int N) {

  if (M < 1024 || N < 1024) {
    _configC(nCore, baseM, baseN, M, N);
  } else {
    nCore = 40;
  }
}

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;
const uint32_t ubDataN = 2;
static ge::graphStatus TilingFunc(gert::TilingContext *context) {
  MatMulSubTilingData tiling;
  auto ascendcPlatform =
      platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  auto shape_x1 = context->GetInputTensor(0)->GetOriginShape();
  auto shape_x2 = context->GetInputTensor(1)->GetOriginShape();
  auto shape_x3 = context->GetInputTensor(2)->GetOriginShape();
  int32_t M = shape_x1.GetDim(0);
  int32_t N = shape_x2.GetDim(1);
  int32_t K = shape_x1.GetDim(1);
  uint32_t x3ShapeType = shape_x3.GetDimNum();
  // SET MATMUL Tiling
  MultiCoreMatmulTiling cubeTiling(ascendcPlatform);
  // TODO:根据M N确定合适的baseM和baseN
  int32_t baseM = 128;
  int32_t baseN = 256;
  uint32_t dataTypeSize =
      GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
  // 暂时仅使用双核
  uint32_t nCore = 4;
  configCores(nCore, baseM, baseN, M, N);
  if (M <= 32 && N <= 32)
    nCore /= 2;
  cubeTiling.SetDim(nCore);
  context->SetBlockDim(nCore / 2); // tilingDIM的1/2
  cubeTiling.SetShape(M, N, K);
  cubeTiling.SetOrgShape(M, N, K);
  if(baseM==128&&baseN==256){
      cubeTiling.SetFixSplit(baseM, baseN, 64);
  }
  else 
      cubeTiling.SetFixSplit(baseM, baseN, -1);
  cubeTiling.SetBias(false);
  cubeTiling.SetBufferSpace(-1, -1, -1);
  //cubeTiling.SetTraverse(MatrixTraverse::FIRSTN);
  if (dataTypeSize == 2) { // half
    cubeTiling.SetAType(TPosition::GM, CubeFormat::ND,
                        matmul_tiling::DataType::DT_FLOAT16);
    cubeTiling.SetBType(TPosition::GM, CubeFormat::ND,
                        matmul_tiling::DataType::DT_FLOAT16);
    cubeTiling.SetCType(TPosition::GM, CubeFormat::ND,
                        matmul_tiling::DataType::DT_FLOAT16);
  } else { // float
    cubeTiling.SetAType(TPosition::GM, CubeFormat::ND,
                        matmul_tiling::DataType::DT_FLOAT);
    cubeTiling.SetBType(TPosition::GM, CubeFormat::ND,
                        matmul_tiling::DataType::DT_FLOAT);
    cubeTiling.SetCType(TPosition::GM, CubeFormat::ND,
                        matmul_tiling::DataType::DT_FLOAT);
  }
  // Get tiling info & set tiling struct
  if (cubeTiling.GetTiling(tiling.cubeTilingData) == -1) {
    return ge::GRAPH_FAILED;
  }
  tiling.set_x3ShapeType(x3ShapeType);

  // set vec
  int32_t rowPerCore = M / nCore;
  int32_t nBigCores = M % nCore;
  tiling.set_bigCoreRows(rowPerCore + 1);
  tiling.set_smallCoreRows(rowPerCore);
  tiling.set_nBigCore(nBigCores);

  // set 2-dim x2 case tiling
  nCore = 20;
  uint64_t ubSize;
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
  uint32_t sizeOfType =
      GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
  uint32_t Ndata = M * N;
  uint32_t inputSize;
  uint32_t inputSizeAligned;
  inputSize = sizeOfType * Ndata;
  inputSizeAligned = (inputSize + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE;
  uint32_t maxBlockPerIter = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataN;
// printf("maxBlockPerIter=%u!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n", maxBlockPerIter);
  uint32_t maxNElemPerIter = maxBlockPerIter * BLOCK_SIZE / sizeOfType;
  uint32_t blockPerCore = inputSizeAligned / BLOCK_SIZE / nCore;
  uint32_t nTailBlocks = (inputSizeAligned / BLOCK_SIZE) % nCore;
  // 计算Acore需要多少次迭代完成全部计算
  uint32_t nElemAcore = (blockPerCore + 1) * BLOCK_SIZE / sizeOfType;
  uint32_t nIterAcore = (blockPerCore + 1) / maxBlockPerIter;
  nIterAcore = (nIterAcore * maxBlockPerIter >= blockPerCore + 1)
                   ? nIterAcore
                   : nIterAcore + 1;
  uint32_t tailNElemA = nElemAcore % maxNElemPerIter;
  // 计算Bcore(无尾块)需要多少次迭代完成全部计算
  uint32_t nElemBcore = blockPerCore * BLOCK_SIZE / sizeOfType;
  uint32_t nIterBcore = blockPerCore / maxBlockPerIter;
  nIterBcore = (nIterBcore * maxBlockPerIter >= blockPerCore) ? nIterBcore
                                                              : nIterBcore + 1;
  uint32_t tailNElemB = nElemBcore % maxNElemPerIter;
  // 设置TILING结构体
  tiling.set_nAcores(nTailBlocks);
  tiling.set_nIterAcore(nIterAcore);
  tiling.set_nIterBcore(nIterBcore);
  tiling.set_tailNElemA(tailNElemA);
  tiling.set_tailNElemB(tailNElemB);
  tiling.set_nElemPerIter(maxNElemPerIter);
  tiling.set_nElemAcore(nElemAcore);
  tiling.set_nElemBcore(nElemBcore);

  // set workspace
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                      context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  size_t userWorkspaceSize = 40*208*512*4;//baseM * baseN * nCore * sizeof(type)
  size_t systemWorkspaceSize =
      static_cast<size_t>(ascendcPlatform.GetLibApiWorkSpaceSize());
  size_t *currentWorkspace = context->GetWorkspaceSizes(1);
  currentWorkspace[0] = userWorkspaceSize + systemWorkspaceSize;
  return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext *context) {
  const gert::Shape *x1_shape = context->GetInputShape(0);
  gert::Shape *y_shape = context->GetOutputShape(0);
  *y_shape = *x1_shape;
  return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class MatMulSub : public OpDef {
public:
  explicit MatMulSub(const char *name) : OpDef(name) {
    this->Input("x1")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Input("x2")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Input("x3")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Output("y")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

    this->SetInferShape(ge::InferShape);

    this->AICore().SetTiling(optiling::TilingFunc);
    this->AICore().AddConfig("ascend910b");
  }
};

OP_ADD(MatMulSub);
} // namespace ops
